import torch
import itertools
from utils.image_pool import ImagePool
from .base_model import BaseModel
from .styleGAN import networks
from utils.util import SSIM,PSNR,normal,AverageMeter
import numpy as np
from  numpy.fft import fft2,fftshift
# from .lightcnn import network as lightcnn
from models import create_model
# import mypymath
# import time
def get_frequency(src):
            res = []
            for i in range(src.shape[0]):
                res.append(fftshift(fft2(src[i,:,:])))
            return torch.tensor(np.array(res))

class styleGANModel(BaseModel):
    """
    This class implements the CycleGAN model, for learning image-to-image translation without paired data.

    The model training requires '--dataset_mode unaligned' dataset.
    By default, it uses a '--netG resnet_9blocks' ResNet generator,
    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
    and a least-square GANs objective ('--gan_mode lsgan').

    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
        
        Returns:
            the modified parser.

        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
        A (source domain), B (target domain).
        Generators: G_A: A -> B; G_B: B -> A.
        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
        Dropout is not used in the original CycleGAN paper.
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.set_defaults(frequency=True)
        # parser.add_argument('--num_classes', type=int, default=3, help='the architecture situation of model')
        # parser.add_argument('--pretrained_model', type=str, default="lightcnn", help='the architecture of pretrained model')
        # parser.add_argument('--pretrained_path', type=str, default="lightcnn", help='the path of pretrained model')
        # parser.add_argument('--pretrained_ndf', type=int, default=128, help='the number of features of output ')
        # parser.add_argument('--pretrained_output_nc', type=int, default=256, help='the number of features of output ')
        # parser.add_argument('--bins', type=int, default=16, help='the number of features of output ')
        parser.add_argument('--cal_mode', type=str, default='linear', help='the mode of caculating')
        parser.add_argument('--D_generation', type=str, default='mergeS12', help='mergeS12 or source')
        # parser.add_argument('--frequency', type=bool, default=True, help='the mode of caculating')
        # if is_train:
        parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
        parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
        parser.add_argument('--lambda_D', type=float, default=10.0, help='weight for Adversarial network loss')
        parser.add_argument('--lambda_fre', type=float, default=0.05, help='weight for Frequency loss')
        parser.add_argument('--lambda_idt', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
            

        return parser

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        # self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
        # self.loss_chage_line = "\n"
        self.opt.frequency = not self.opt.lambda_fre == 0
        if self.isTrain:
            self.loss_names += ["mean_G_S01","mean_G_S10",
                            "mean_G_S02","mean_G_S20",
                            "mean_G_S0D","mean_G_SD0",
                            "mean_F_S01","mean_F_S10",
                            "mean_F_S02","mean_F_S20",
                            "mean_F_S0D","mean_F_SD0",
                            "mean_F_S12D","mean_G",
                            ]
            self.loss_names += ["mean_D_S0","mean_D_S1","mean_D_S2","mean_D_SD"]

            self.count_mean_G=AverageMeter()
            self.count_mean_G_S10=AverageMeter()
            self.count_mean_G_S01=AverageMeter()
            self.count_mean_G_S20=AverageMeter()
            self.count_mean_G_S02=AverageMeter()
            self.count_mean_G_SD0=AverageMeter()
            self.count_mean_G_S0D=AverageMeter()
            self.count_mean_D_S0 = AverageMeter()
            self.count_mean_D_S1 = AverageMeter()
            self.count_mean_D_S2 = AverageMeter()
            self.count_mean_D_SD = AverageMeter()
    # else:
        self.loss_names += ["S1_SSIM","S01_SSIM","S2_SSIM","S02_SSIM"]
        self.loss_names += ["S1_PSNR","S01_PSNR","S2_PSNR","S02_PSNR",]
        self.loss_names += ["mean_S1_PSNR","mean_S01_PSNR","mean_S2_PSNR","mean_S02_PSNR"]
        self.count_mean_S01_PSNR=AverageMeter()
        self.count_mean_S02_PSNR=AverageMeter()
        self.count_mean_S1_PSNR=AverageMeter()
        self.count_mean_S2_PSNR=AverageMeter()
        self.count_mean_S0D_PSNR=AverageMeter()
        self.val_best = 0
        self.last_val_best = 0
        
        self.loss_names += ["mean_S1_SSIM","mean_S01_SSIM","mean_S2_SSIM","mean_S02_SSIM"]
        
        self.count_mean_S1_SSIM=AverageMeter()
        self.count_mean_S2_SSIM=AverageMeter()
        self.count_mean_S0D_SSIM=AverageMeter()
        self.count_mean_S01_SSIM=AverageMeter()
        self.count_mean_S02_SSIM=AverageMeter()
        
        if self.opt.D_generation == "source":
            # self.loss_names += ["SD_PSNR","S0D_PSNR","SD_SSIM","S0D_SSIM",]
            self.loss_names += ["mean_SD_SSIM","mean_S0D_SSIM"]
            self.count_mean_SD_SSIM=AverageMeter()
            
            # if not self.isTrain:
            self.loss_names += ["mean_SD_PSNR","mean_S0D_PSNR"]
            self.count_mean_SD_PSNR=AverageMeter()
            
            
        elif self.opt.D_generation == "mergeS12":
            # self.loss_names += ["S12D_PSNR","S0_12D_PSNR","S12D_SSIM","S0_12D_SSIM"]
            self.loss_names += ["mean_S12D_SSIM","mean_S0D_SSIM"]
            self.count_mean_S12D_SSIM=AverageMeter()
            # if not self.isTrain:
            self.loss_names += ["mean_S12D_PSNR","mean_S0D_PSNR",]
            self.count_mean_S12D_PSNR=AverageMeter()
            
            
        if self.opt.frequency:
            self.count_mean_F_S01=AverageMeter()
            self.count_mean_F_S10=AverageMeter()
            self.count_mean_F_S02=AverageMeter()
            self.count_mean_F_S20=AverageMeter()
            self.count_mean_F_S0D=AverageMeter()
            self.count_mean_F_SD0=AverageMeter()
            self.count_mean_F_S12D=AverageMeter()
        
        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        visual_names_S1A = []
        visual_names_S1B = []
        visual_names_S1A += ['real_S0', 'fake_S01', 'rec_S10']
        # visual_names_S1A += ['real_S0_h', 'fake_S01_h', 'rec_S10_h']
        visual_names_S1A += ['fake_S02', 'rec_S20']
        # visual_names_S1A += [ 'fake_S02_h', 'rec_S20_h']
        visual_names_S1A += [ 'fake_S0D', 'rec_SD0']
        # visual_names_S1A += [ 'fake_S0D_h', 'rec_SD0_h']
        visual_names_S1A += ['real_S1', 'real_S2', 'real_SD', 'fake_S12D']
        # visual_names_S1B += [, 'real_S12D',]
        # visual_names_S1A += ['real_S2_h', 'real_SD_h', 'real_S1_h']
        
        if self.isTrain:
            visual_names_S1B += ['fake_S10', 'rec_S01']
            # visual_names_S1B += ['fake_S10_h', 'rec_S01_h']
            visual_names_S1B += ['fake_S20', 'rec_S02']
            # visual_names_S1B += ['fake_S20_h', 'rec_S02_h']
            visual_names_S1B += ['fake_SD0', 'rec_S0D']
            # visual_names_S1B += ['fake_SD0_h', 'rec_S0D_h']
        
        self.visual_names = visual_names_S1A + visual_names_S1B  # combine visualizations for A and B
        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        if self.isTrain:
            # self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
            self.model_names = ['G_S01', 'G_S10', 
                                'G_S02', 'G_S20', 
                                'D_S0', 'D_S1', 'D_S2','D_SD',]
            if self.opt.D_generation == "source":
                self.model_names += ['G_S0D', 'G_SD0',]
        else:  # during test time, only load Gs
            self.model_names = ['G_S01', 'G_S10','G_S02', 'G_S20', ]
            if self.opt.D_generation == "source":
                self.model_names += ['G_S0D', 'G_SD0',]

        # define networks (both Generators and discriminators)
        # The naming is different from those used in the paper.
        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
        self.netG_S01 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
        
        self.netG_S10 = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
        
        
        self.netG_S02 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
        
        self.netG_S20 = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
        
        if self.opt.D_generation == 'source':
            self.netG_S0D = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
            
            self.netG_SD0 = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)


        if self.isTrain:  # define discriminators
            self.netD_S1 = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)

            self.netD_S0 = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
            
            self.netD_S2 = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
            
            self.netD_SD = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
            # self.netD_light.define_lightcnn(opt)

            # self.netD_S02 = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
            #                                 opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
        self.criterionPredD = torch.nn.L1Loss()
        
        self.opt.lambda_idt = self.opt.lambda_idt if self.opt.lambda_idt > 0 else 0
            # define loss functions
        self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
        self.criterionCycle = torch.nn.L1Loss()
        self.criterionIdt = torch.nn.L1Loss()
        
        if self.isTrain:
            if self.opt.continue_train:
                self.epoch = self.opt.epoch_count
            # self.nepoch = self.opt.n_epochs + self.opt.n_epochs_decay
            # self.fac = self.epoch/(self.nepoch+1)
            if opt.lambda_idt > 0.0:  # only works when input and output images have the same number of channels
                assert(opt.input_nc == opt.output_nc)
            self.fake_S0_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fake_S1_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fake_S2_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            self.fake_SD_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # self.fake_S2_S0_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            
        # if self.isTrain:
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            if self.opt.D_generation == 'source':
                self.optimizer_G = torch.optim.Adam(itertools.chain(
                                                                self.netG_S01.parameters(),
                                                                self.netG_S10.parameters(),
                                                                self.netG_S02.parameters(),
                                                                self.netG_S20.parameters(),
                                                                self.netG_S0D.parameters(),
                                                                self.netG_SD0.parameters(),
                                                                ),
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))
            elif self.opt.D_generation == 'mergeS12':
                self.optimizer_G = torch.optim.Adam(itertools.chain(
                                                                self.netG_S01.parameters(),
                                                                self.netG_S10.parameters(),
                                                                self.netG_S02.parameters(),
                                                                self.netG_S20.parameters(),
                                                                ),
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))

            self.optimizer_D = torch.optim.Adam(itertools.chain(
                                                                self.netD_S0.parameters(),
                                                                self.netD_S1.parameters(),
                                                                self.netD_S2.parameters(),
                                                                self.netD_SD.parameters(),
                                                                ),
                                                lr=opt.lr, betas=(opt.beta1, 0.999))

            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

            

    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        # AtoB = self.opt.direction == 'AtoB'
        # self.real_A = input['A' if AtoB else 'B'].to(self.device)
        # self.real_B = input['B' if AtoB else 'A']self.device)
        # self.image_paths = input['A_paths' if AtoB else 'B_paths']
        if "isTrain" in input.keys():
            self.isTrain = input["isTrain"][0].item()
        
        self.real_S0 = input['S0'].to(self.device)
        self.image_paths = input['S0_paths']
        # if self.isTrain:
        self.real_S1 = input['S1'].to(self.device)
        self.real_S2 = input['S2'].to(self.device)
        self.real_SD = input['SD'].to(self.device)
        
        # self.real_S0 = normal(self.real_S0,1,1e-5)
        # self.real_S1 = normal(self.real_S1,1,1e-5)
        # self.real_S2 = normal(self.real_S2,1,1e-5)
        # self.real_SD = normal(self.real_SD,1,1e-5)
            # self.real_S2 = input['S2'].to(self.device)
            
        # self.real_S12D = torch.sqrt(self.real_S1**2 + self.real_S2**2)/(self.real_S0+1e-5)
        # self.real_S12D = normal(self.real_S12D,1,1e-5)
                
        if self.opt.frequency:
            self.real_S0_F = input['S0_F'][:,0,:,:]
            self.real_S1_F = input['S1_F'][:,0,:,:]
            self.real_S2_F = input['S2_F'][:,0,:,:]
            self.real_SD_F = input['SD_F'][:,0,:,:]
        
        
            

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        # self.fake_B = self.netG_A(self.real_A)  # G_A(A)
        # self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
        # self.fake_A = self.netG_B(self.real_B)  # G_B(B)
        # self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))
        self.fake_S01 = self.netG_S01(self.real_S0)
        self.rec_S10 = self.netG_S10(self.fake_S01)
        
        self.fake_S02 = self.netG_S02(self.real_S0)
        self.rec_S20 = self.netG_S20(self.fake_S02)
        
        if self.opt.D_generation == "source":
            self.fake_S0D = self.netG_S0D(self.real_S0)
            self.rec_SD0 = self.netG_SD0(self.fake_S0D)
        
        if self.isTrain:
            
            self.fake_S10 = self.netG_S10(self.real_S1)
            self.rec_S01 = self.netG_S01(self.fake_S10)
            
            self.fake_S20 = self.netG_S20(self.real_S2)
            self.rec_S02 = self.netG_S02(self.fake_S20)
            
            if self.opt.D_generation == "source":
                self.fake_SD0 = self.netG_SD0(self.real_SD)
                self.rec_S0D = self.netG_S0D(self.fake_SD0)
            
        # self.fake_S01 = normal(self.fake_S01,1,1e-5)
        # self.fake_S02 = normal(self.fake_S02,1,1e-5)
        # x1 = self.fake_S01/self.real_S0
        # x2 = self.fake_S02/self.real_S0
        # x1 = normal(x1,1,-1)
        # x2 = normal(x2,1,-1)
        if self.opt.D_generation == "mergeS12":
            self.fake_S12D = torch.sqrt(self.fake_S01**2 + self.fake_S02**2)
        # self.fake_S12D = normal(self.fake_S12D,1,1e-5)
        
    def backward_D_basic(self, netD, real, fake):
        """Calculate GAN loss for the discriminator

        Parameters:
            netD (network)      -- the discriminator D
            real (tensor array) -- real images
            fake (tensor array) -- images generated by a generator

        Return the discriminator loss.
        We also call loss_D.backward() to calculate the gradients.
        """
        # Real
        pred_real = netD(real)
        loss_D_real = self.criterionGAN(pred_real, True)
        # Fake
        pred_fake = netD(fake.detach())
        loss_D_fake = self.criterionGAN(pred_fake, False)
        # Combined loss and calculate gradients
        loss_D = (loss_D_real + loss_D_fake) * 0.5
        if self.isTrain:
            loss_D.backward()
        return loss_D

    def backward_D_S0(self):
        """Calculate GAN loss for discriminator D_S0"""
        if self.opt.D_generation == "source":
            fake_S0 = torch.cat([self.fake_S10,self.fake_S20,self.fake_SD0])
        elif self.opt.D_generation == "mergeS12":
            fake_S0 = torch.cat([self.fake_S10,self.fake_S20])
        fake_S0 = self.fake_S0_pool.query(fake_S0)
        self.loss_D_S0 = self.backward_D_basic(self.netD_S0, self.real_S0, fake_S0)
        self.count_mean_D_S0.update(self.loss_D_S0)
        self.loss_mean_D_S0 = self.count_mean_D_S0.avg
        
    def backward_D_S1(self):
        """Calculate GAN loss for discriminator D_S1"""
        fake_S01 = self.fake_S1_pool.query(self.fake_S01)
        self.loss_D_S1 = self.backward_D_basic(self.netD_S1, self.real_S1, fake_S01)
        self.count_mean_D_S1.update(self.loss_D_S1)
        self.loss_mean_D_S1 = self.count_mean_D_S1.avg

    
    def backward_D_S2(self):
        """Calculate GAN loss for discriminator D_S1"""
        fake_S02 = self.fake_S2_pool.query(self.fake_S02)
        self.loss_D_S2 = self.backward_D_basic(self.netD_S2, self.real_S2, fake_S02)
        self.count_mean_D_S2.update(self.loss_D_S2)
        self.loss_mean_D_S2 = self.count_mean_D_S2.avg
        
        
    def backward_D_SD(self):
        """Calculate GAN loss for discriminator D_S1"""
        if self.opt.D_generation == "source":
            fake_S0D = self.fake_S1_pool.query(self.fake_S0D)
        elif self.opt.D_generation == "mergeS12":
            fake_S0D = self.fake_S1_pool.query(self.fake_S12D)
            
        # fake_S0D = self.fake_S1_pool.query(torch.cat([self.fake_S0D,self.fake_S12D]))
        self.loss_D_SD = self.backward_D_basic(self.netD_SD, self.real_SD, fake_S0D)
        self.count_mean_D_SD.update(self.loss_D_SD)
        self.loss_mean_D_SD = self.count_mean_D_SD.avg
    
    def set_Frequency_S01(self,loss1,loss2):
        self.loss_F_S01 = (loss1 + loss2).to(self.device)
    
    def set_Frequency_S10(self,loss1,loss2):
        self.loss_F_S10 = (loss1 + loss2).to(self.device)
        
    def set_Frequency_S02(self,loss1,loss2):
        self.loss_F_S02 = (loss1 + loss2).to(self.device)
        
    def set_Frequency_S20(self,loss1,loss2):
        self.loss_F_S20 = (loss1 + loss2).to(self.device)
    
    def set_Frequency_S0D(self,loss1,loss2):
        self.loss_F_S0D = (loss1 + loss2).to(self.device)
    
    def set_Frequency_SD0(self,loss1,loss2):
        self.loss_F_SD0 = (loss1 + loss2).to(self.device)
    
    def backward_Frequency(self,X):
        
        # starttime1 = time.time()
        fake_S0X_F = get_frequency(getattr(self,'fake_S0'+X).contiguous().cpu().detach().numpy()[:,0,:,:])
        rec_SX0_F = get_frequency(getattr(self,'rec_S'+X+'0').contiguous().cpu().detach().numpy()[:,0,:,:])
        if self.isTrain:
            rec_S0X_F = get_frequency(getattr(self,'rec_S0'+X).contiguous().cpu().detach().numpy()[:,0,:,:])
            fake_SX0_F = get_frequency(getattr(self,'fake_S'+X+'0').contiguous().cpu().detach().numpy()[:,0,:,:])
            getattr(self,'set_Frequency_S0'+X)(self.criterionPredD(fake_S0X_F,getattr(self,'real_S'+X+'_F')),self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')))
            getattr(self,'set_Frequency_S'+X+'0')(self.criterionPredD(fake_SX0_F,getattr(self,'real_S0_F')),self.criterionPredD(rec_SX0_F,getattr(self,'real_S0_F')))
        # getattr(self,'set_Frequency_rec_S0'+'X')(self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')),self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')))
        # getattr(self,'set_Frequency_rec_S'+X+'0')(self.criterionPredD(rec_SX0_F,getattr(self,'real_S0_F')),self.criterionPredD(rec_S0X_F,getattr(self,'real_S'+X+'_F')))
        # getattr(self,'loss_Frequency'+) = self.loss_Frequency_fake_S0 + self.loss_Frequency_fake_S1 + self.loss_Frequency_rec_S0 + self.loss_Frequency_rec_S1
        # getattr(self,count_Frequency).update(self.loss_Frequency.item())
        # getattr(self,loss_mean_Frequency) = self.count_Frequency.avg

    def cal_param_G_S01(self):
        self.loss_G_S01 = torch.tensor(float(0)).to(self.device)
        idt_S01 = self.netG_S01(self.real_S1)
        loss_idt_S01 = self.criterionIdt(idt_S01, self.real_S1) * self.opt.lambda_B * self.opt.lambda_idt
        loss_cycle_S01 = (self.criterionCycle(self.rec_S01, self.real_S1) )* self.opt.lambda_B
        
        self.loss_G_S01 += loss_idt_S01 +loss_cycle_S01
        if self.opt.frequency:
            self.loss_G_S01 += self.loss_F_S01 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S01 += self.criterionGAN(self.netD_S1(self.fake_S01), True)*self.opt.lambda_D
            # self.loss_G_S01 += self.criterionCycle(self.fake_S01, self.real_S1)
        pass
    
    def cal_param_G_S10(self):
        self.loss_G_S10 = torch.tensor(float(0)).to(self.device)
        idt_S10= self.netG_S10(self.real_S0)
        loss_idt_S10 = self.criterionIdt(idt_S10, self.real_S0) * self.opt.lambda_A * self.opt.lambda_idt
        loss_cycle_S10 = self.criterionCycle(self.rec_S10, self.real_S0) * self.opt.lambda_A
        
        self.loss_G_S10 += loss_cycle_S10 + loss_idt_S10
        if self.opt.frequency:
            self.loss_G_S10 += self.loss_F_S10 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S10 += self.criterionGAN(self.netD_S0(self.fake_S10), True)*self.opt.lambda_D
            # self.loss_G_S10 += self.criterionCycle(self.fake_S10, self.real_S0)
            # self.loss_G_S10.backward()
        pass
    
    def cal_param_G_S02(self):
        self.loss_G_S02 = torch.tensor(float(0)).to(self.device)
        idt_S02 = self.netG_S02(self.real_S2)
        loss_idt_S02 = self.criterionIdt(idt_S02, self.real_S2) * self.opt.lambda_B * self.opt.lambda_idt
        loss_cycle_S02 = self.criterionCycle(self.rec_S02, self.real_S2) * self.opt.lambda_B
        
        self.loss_G_S02 += loss_idt_S02 +loss_cycle_S02
        if self.opt.frequency:
            self.loss_G_S02 += self.loss_F_S02 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S02 += self.criterionGAN(self.netD_S2(self.fake_S02), True)*self.opt.lambda_D
            # self.loss_G_S02 += self.criterionCycle(self.fake_S02, self.real_S2)
            # self.loss_G_S02.backward()
        pass
    
    def cal_param_G_S20(self):
        self.loss_G_S20 = torch.tensor(float(0)).to(self.device)
        idt_S20= self.netG_S20(self.real_S0)
        loss_idt_S20 = self.criterionIdt(idt_S20, self.real_S0) * self.opt.lambda_A * self.opt.lambda_idt
        loss_cycle_S20 = self.criterionCycle(self.rec_S20, self.real_S0) * self.opt.lambda_A
        
        self.loss_G_S20 += loss_cycle_S20 + loss_idt_S20
        if self.opt.frequency:
            self.loss_G_S20 += self.loss_F_S20 * self.opt.lambda_fre
        if self.isTrain:
            self.loss_G_S20 += self.criterionGAN(self.netD_S0(self.fake_S20), True)*self.opt.lambda_D
            # self.loss_G_S20 += self.criterionCycle(self.fake_S20, self.real_S0)
            # self.loss_G_S20.backward()
        pass
    
    def cal_param_G_S0D(self):
        self.loss_G_S0D = torch.tensor(float(0)).to(self.device)
        if self.opt.D_generation == "source":
            idt_S0D = self.netG_S0D(self.real_SD)
            loss_idt_S0D = self.criterionIdt(idt_S0D, self.real_SD) * self.opt.lambda_B * self.opt.lambda_idt
            loss_cycle_S0D = self.criterionCycle(self.rec_S0D, self.real_SD) * self.opt.lambda_B
            self.loss_G_S0D += loss_idt_S0D +loss_cycle_S0D
            if self.opt.frequency:
                self.loss_G_S0D += self.loss_F_S0D * self.opt.lambda_fre
        elif self.opt.D_generation == "mergeS12":
            if self.opt.frequency:
                self.loss_G_S0D += self.loss_F_S12D * self.opt.lambda_fre
        if self.isTrain:
            if self.opt.D_generation == "source":
                self.loss_G_S0D += self.criterionGAN(self.netD_SD(self.fake_S0D), True)*self.opt.lambda_D
                # self.loss_G_S0D += self.criterionCycle(self.fake_S0D, self.real_SD)
            elif self.opt.D_generation == "mergeS12":
                self.loss_G_S0D += self.criterionGAN(self.netD_SD(self.fake_S12D), True)*self.opt.lambda_D
                # self.loss_G_S0D += self.criterionCycle(self.fake_S12D, self.real_SD)
            # self.loss_G_S0D.backward()
        pass
    
    def cal_param_G_SD0(self):
        self.loss_G_SD0 = torch.tensor(float(0)).to(self.device)
        if self.opt.D_generation == "source":
            idt_SD0= self.netG_SD0(self.real_S0)
            loss_idt_SD0 = self.criterionIdt(idt_SD0, self.real_S0) * self.opt.lambda_A * self.opt.lambda_idt
            loss_cycle_SD0 = self.criterionCycle(self.rec_SD0, self.real_S0) * self.opt.lambda_A
            
            self.loss_G_SD0 += loss_cycle_SD0 + loss_idt_SD0
            if self.opt.frequency:
                self.loss_G_SD0 += self.loss_F_SD0 * self.opt.lambda_fre
            if self.isTrain:
                self.loss_G_SD0 += self.criterionGAN(self.netD_S0(self.fake_SD0), True)*self.opt.lambda_D
                # self.loss_G_SD0 += self.criterionCycle(self.fake_SD0, self.real_S0)
            # self.loss_G_SD0.backward()
        pass
    
    def calculate_G_parameters(self):
        # if self.isTrain:
        if self.opt.frequency:
            self.backward_Frequency('1')
            self.backward_Frequency('2')
            if self.opt.D_generation == "source":
                self.backward_Frequency('D')
            elif self.opt.D_generation == "mergeS12":
                fake_S12D_F = get_frequency(self.fake_S12D.contiguous().cpu().detach().numpy()[:,0,:,:])
                self.loss_F_S12D = self.criterionPredD(fake_S12D_F , self.real_SD_F).to(self.device)
            if self.isTrain:
                self.count_mean_F_S01.update(self.loss_F_S01.contiguous().detach().cpu())
                self.count_mean_F_S10.update(self.loss_F_S10.contiguous().detach().cpu())
                self.count_mean_F_S02.update(self.loss_F_S02.contiguous().detach().cpu())
                self.count_mean_F_S20.update(self.loss_F_S20.contiguous().detach().cpu())
                
                if self.opt.D_generation == "source":
                    self.count_mean_F_S0D.update(self.loss_F_S0D.contiguous().detach().cpu())
                    self.count_mean_F_SD0.update(self.loss_F_SD0.contiguous().detach().cpu())
                elif self.opt.D_generation == "mergeS12":
                    self.count_mean_F_S12D.update(self.loss_F_S12D.contiguous().detach().cpu())
                
                self.loss_mean_F_S01 = self.count_mean_F_S01.avg
                self.loss_mean_F_S10 = self.count_mean_F_S10.avg
                self.loss_mean_F_S02 = self.count_mean_F_S02.avg
                self.loss_mean_F_S20 = self.count_mean_F_S20.avg
                
                if self.opt.D_generation == "source":
                    self.loss_mean_F_S0D = self.count_mean_F_S0D.avg
                    self.loss_mean_F_SD0 = self.count_mean_F_SD0.avg
                elif self.opt.D_generation == "mergeS12":
                    self.loss_mean_F_S12D = self.count_mean_F_S12D.avg
        
        if self.isTrain:
            self.cal_param_G_S01()
            self.cal_param_G_S10()
            self.cal_param_G_S02()
            self.cal_param_G_S20()
            self.cal_param_G_S0D()
            self.cal_param_G_SD0()
            
            self.loss_G = self.loss_G_S01 + self.loss_G_S10
            self.loss_G += self.loss_G_S02 + self.loss_G_S20
            self.loss_G += self.loss_G_S0D + self.loss_G_SD0
            self.loss_G.backward()
 
            self.count_mean_G.update(self.loss_G.contiguous().detach().cpu())
            self.count_mean_G_S01.update(self.loss_G_S01.contiguous().detach().cpu())
            self.count_mean_G_S10.update(self.loss_G_S10.contiguous().detach().cpu())
            self.count_mean_G_S02.update(self.loss_G_S02.contiguous().detach().cpu())
            self.count_mean_G_S20.update(self.loss_G_S20.contiguous().detach().cpu())
            self.count_mean_G_S0D.update(self.loss_G_S0D.contiguous().detach().cpu())
            self.count_mean_G_SD0.update(self.loss_G_SD0.contiguous().detach().cpu())
            
            self.loss_mean_G=self.count_mean_G.avg
            self.loss_mean_G_S01=self.count_mean_G_S01.avg
            self.loss_mean_G_S10=self.count_mean_G_S10.avg
            self.loss_mean_G_S02=self.count_mean_G_S02.avg
            self.loss_mean_G_S20=self.count_mean_G_S20.avg
            self.loss_mean_G_S0D=self.count_mean_G_S0D.avg
            self.loss_mean_G_SD0=self.count_mean_G_SD0.avg
            self.cal_score()
            
        pass
    
    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # forward
        self.forward()      # compute fake images and reconstruction images.
        
        # G_A and G_B
        self.set_requires_grad([self.netD_S0, self.netD_S1,self.netD_S2, self.netD_SD,], False)  # Ds require no gradients when optimizing Gs
        self.optimizer_G.zero_grad()
        self.calculate_G_parameters()
        self.optimizer_G.step()

        self.set_requires_grad([self.netD_S0, self.netD_S1,self.netD_S2, self.netD_SD,], True)
        
        self.optimizer_D.zero_grad()
        self.backward_D_S0()
        self.backward_D_S1()
        self.backward_D_S2()
        self.backward_D_SD()
        self.optimizer_D.step()
        
    def cal_score(self):
        if not self.isTrain:
            self.calculate_G_parameters()
        # self.real_S0 /= self.real_S0.max()
        # self.rec_S10 /= self.rec_S10.max()
        # self.rec_S20 /= self.rec_S20.max()
        # self.rec_SD0 /= self.rec_SD0.max()
        
        
        self.real_S10 = normal(self.rec_S10,1,-1)
        self.real_S20 = normal(self.rec_S20,1,-1)
        
        # self.real_S1 *= self.real_S0
        # self.fake_S01 *= self.real_S0
        # self.real_S2 *= self.real_S0
        # self.fake_S02 *= self.real_S0
        
        # self.fake_S01 = normal(self.fake_S01,self.real_S1.max(),self.real_S1.min())
        # self.fake_S02 = normal(self.fake_S02,self.real_S2.max(),self.real_S2.min())
        
        self.real_S0 = normal(self.real_S0,1,-1)
        self.real_S1 = normal(self.real_S1,1,-1)
        self.fake_S01 = normal(self.fake_S01,1,-1)
        self.real_S2 = normal(self.real_S2,1,-1)
        self.fake_S02 = normal(self.fake_S02,1,-1)
        self.real_SD = normal(self.real_SD,1,-1)
        
        if not self.isTrain:
            self.loss_S1_PSNR = PSNR(self.fake_S01,self.real_S1)
            self.loss_S2_PSNR = PSNR(self.fake_S02,self.real_S2)
            self.loss_S01_PSNR = PSNR(self.real_S0,self.real_S1)
            self.loss_S02_PSNR = PSNR(self.real_S0,self.real_S2)
            self.loss_S0D_PSNR = PSNR(self.real_S0,self.real_SD)
            self.count_mean_S1_PSNR.update(self.loss_S1_PSNR)
            self.count_mean_S2_PSNR.update(self.loss_S2_PSNR)
            self.count_mean_S01_PSNR.update(self.loss_S01_PSNR)
            self.count_mean_S02_PSNR.update(self.loss_S02_PSNR)
            self.count_mean_S0D_PSNR.update(self.loss_S0D_PSNR)
            self.loss_mean_S1_PSNR = self.count_mean_S1_PSNR.avg
            self.loss_mean_S2_PSNR = self.count_mean_S2_PSNR.avg
            self.loss_mean_S01_PSNR = self.count_mean_S01_PSNR.avg
            self.loss_mean_S02_PSNR = self.count_mean_S02_PSNR.avg
            self.loss_mean_S0D_PSNR = self.count_mean_S0D_PSNR.avg

        self.loss_S1_SSIM = SSIM(self.fake_S01,self.real_S1)
        self.loss_S2_SSIM = SSIM(self.fake_S02,self.real_S2)
        self.loss_S01_SSIM = SSIM(self.real_S0,self.real_S1)
        self.loss_S02_SSIM = SSIM(self.real_S0,self.real_S2)
        self.loss_S0D_SSIM = SSIM(self.real_S0,self.real_SD)
        
        self.count_mean_S1_SSIM.update(self.loss_S1_SSIM)
        self.count_mean_S2_SSIM.update(self.loss_S2_SSIM)
        self.count_mean_S01_SSIM.update(self.loss_S01_SSIM)
        self.count_mean_S02_SSIM.update(self.loss_S02_SSIM)
        self.count_mean_S0D_SSIM.update(self.loss_S0D_SSIM)
        
        
        self.loss_mean_S1_SSIM = self.count_mean_S1_SSIM.avg
        self.loss_mean_S2_SSIM = self.count_mean_S2_SSIM.avg
        self.loss_mean_S01_SSIM = self.count_mean_S01_SSIM.avg
        self.loss_mean_S02_SSIM = self.count_mean_S02_SSIM.avg
        self.loss_mean_S0D_SSIM = self.count_mean_S0D_SSIM.avg

        if self.opt.D_generation == "source":
                self.rec_SD0 = normal(self.rec_SD0,1,-1)
                self.fake_S0D = normal(self.fake_S0D,1,-1)
                if not self.isTrain:
                    self.loss_SD_PSNR = PSNR(self.fake_S0D,self.real_SD)
                    self.count_mean_SD_PSNR.update(self.loss_SD_PSNR)
                    self.loss_mean_SD_PSNR = self.count_mean_SD_PSNR.avg
                
                self.loss_SD_SSIM = SSIM(self.fake_S0D,self.real_SD)
                self.count_mean_SD_SSIM.update(self.loss_SD_SSIM)
                self.loss_mean_SD_SSIM = self.count_mean_SD_SSIM.avg
                
                
        elif self.opt.D_generation == "mergeS12":
            self.fake_S12D = normal(self.fake_S12D,1,-1)
            
            if not self.isTrain:
                self.loss_S12D_PSNR = PSNR(self.fake_S12D,self.real_SD)
                self.count_mean_S12D_PSNR.update(self.loss_S12D_PSNR)
                self.loss_mean_S12D_PSNR = self.count_mean_S12D_PSNR.avg
            
            self.loss_S12D_SSIM = SSIM(self.fake_S12D,self.real_SD)
            self.count_mean_S12D_SSIM.update(self.loss_S12D_SSIM)
            self.loss_mean_S12D_SSIM = self.count_mean_S12D_SSIM.avg
        pass

        pass
    def iter_end(self):
        pass
    def epoch_end(self):
        if self.isTrain:
            self.count_mean_G.reset()
            self.count_mean_G_S10.reset()
            self.count_mean_G_S01.reset()
            self.count_mean_G_S20.reset()
            self.count_mean_G_S02.reset()
            self.count_mean_G_SD0.reset()
            self.count_mean_G_S0D.reset()
            self.count_mean_D_S0.reset()
            self.count_mean_D_S1.reset()
            self.count_mean_D_S2.reset()
            self.count_mean_D_SD.reset()
        else:
            self.count_mean_S1_PSNR.reset()
            self.count_mean_S2_PSNR.reset()
            self.count_mean_S01_PSNR.reset()
            self.count_mean_S02_PSNR.reset()
            self.count_mean_S0D_PSNR.reset()
            
            
            
        self.count_mean_S1_SSIM.reset()
        self.count_mean_S2_SSIM.reset()
        self.count_mean_S01_SSIM.reset()
        self.count_mean_S02_SSIM.reset()
        self.count_mean_S0D_SSIM.reset()
        
        if self.opt.D_generation == "source":
            if not self.isTrain:  
                self.count_mean_SD_PSNR.reset()
            self.count_mean_SD_SSIM.reset()
        elif self.opt.D_generation == "mergeS12":
            if not self.isTrain:
                self.count_mean_S12D_PSNR.reset()
            self.count_mean_S12D_SSIM.reset()
                         
        if self.opt.frequency:
            self.count_mean_F_S01.reset()
            self.count_mean_F_S10.reset()
            self.count_mean_F_S02.reset()
            self.count_mean_F_S20.reset()
            self.count_mean_F_S0D.reset()
            self.count_mean_F_SD0.reset()
            self.count_mean_F_S12D.reset()
    def val_save(self):
        self.val_best = (self.count_mean_S1_SSIM.avg + self.count_mean_S2_SSIM.avg + self.count_mean_SD_SSIM.avg)
        if self.last_val_best < self.val_best:
            self.last_val_best = self.val_best
            return True
        return False
        
        # if self.opt.isTrain:
        #     self.epoch+=1
            # self.fac = self.epoch/(self.nepoch+1)
